import pymysql
from contextlib import contextmanager
from typing import Dict, Optional, Iterator, Any
import yaml
import sys

#数据库配置文件参照db_config_example
def load_config():
    with open('db_config.yaml', 'r', encoding='utf-8') as f:
        config = yaml.safe_load(f)
        return config

class DatabaseConnection:

    def __init__(
        self,
        host: str = '',
        user: str = '',
        password: str= '',
        database: str = '',
        port: int = 3306,
        charset: str = 'utf8mb4',
        autocommit: bool = False
    ):
        """
        初始化数据库连接参数

    host: 数据库主机地址
    user: 数据库用户名
    password: 数据库密码
    database: 数据库名称
    port: 数据库端口，默认为3306
    charset: 字符集，默认为utf8mb4
    autocommit: 是否自动提交事务，默认为False
        """
        db_config=load_config(),
        print(db_config)
        self.host = db_config[0]['mysql']['host']
        self.user = db_config[0]['mysql']['user']
        self.password = db_config[0]['mysql']['password']
        self.database = db_config[0]['mysql']['database']
        self.port = db_config[0]['mysql']['port']
        self.charset = db_config[0]['mysql']['charset']
        self.autocommit = autocommit
        self.connection: Optional[pymysql.Connection] = None

    def connect(self) -> pymysql.Connection:
        """
        建立数据库连接

        :return: 数据库连接对象
        :raises: pymysql.Error 如果连接失败
        """
        try:
            self.connection = pymysql.connect(
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.password,
                database=self.database,
                charset=self.charset,
                autocommit=self.autocommit,
                cursorclass=pymysql.cursors.DictCursor
            )
            return self.connection
        except pymysql.Error as e:
            raise ConnectionError(f"数据库连接失败: {e}") from e

    def disconnect(self):
        """断开数据库连接"""
        if self.connection and self.connection.open:
            try:
                self.connection.close()
            except pymysql.Error:
                pass  # 忽略关闭时的异常
        self.connection = None

    def is_connected(self) -> bool:
        """检查是否已连接数据库"""
        return self.connection is not None and self.connection.open

    @contextmanager
    def connection_context(self) -> Iterator[pymysql.Connection]:
        """
        数据库连接上下文管理器

        用法:
        with db.connection_context() as conn:
            with conn.cursor() as cursor:
                cursor.execute("SELECT ...")
                results = cursor.fetchall()
        """
        conn = self.connect()
        try:
            yield conn
        finally:
            self.disconnect()

    @contextmanager
    def cursor_context(self) -> Iterator[pymysql.cursors.Cursor]:
        """
        游标上下文管理器（自动处理连接和游标）

        用法:
        with db.cursor_context() as cursor:
            cursor.execute("SELECT ...")
            results = cursor.fetchall()
        """
        conn = self.connect()
        try:
            with conn.cursor() as cursor:
                yield cursor
                # 如果不是自动提交模式，需要手动提交事务
                if not self.autocommit:
                    conn.commit()
        except pymysql.Error:
            # 发生异常时回滚事务
            if conn.open and not self.autocommit:
                conn.rollback()
            raise
        finally:
            self.disconnect()

    def test_connection(self) -> bool:
        """测试数据库连接是否正常"""
        try:
            with self.cursor_context() as cursor:
                cursor.execute("SELECT 1")
                result = cursor.fetchone()
                return result is not None and result.get('1') == 1
        except pymysql.Error:
            return False

    def get_connection_info(self) -> Dict[str, Any]:
        """获取数据库连接信息"""
        return {
            'host': self.host,
            'port': self.port,
            'user': self.user,
            'database': self.database,
            'charset': self.charset,
            'autocommit': self.autocommit,
            'connected': self.is_connected()
        }


# 使用示例，以下仅供测试使用
if __name__ == "__main__":
    # 1. 创建数据库连接配置
    # 2. 创建数据库连接对象
    db = DatabaseConnection()

    # 3. 测试连接
    print("测试数据库连接...", end=' ')
    if db.test_connection():
        print("成功!")
        print("连接信息:", db.get_connection_info())
    else:
        print("失败!")
        exit(1)

    # 4. 使用连接上下文
    print("\n使用连接上下文:")
    with db.connection_context() as conn:
        with conn.cursor() as cursor:
            cursor.execute("SHOW TABLES")
            tables = cursor.fetchall()
            print("数据库中的表:")
            for table in tables:
                print(table[f"Tables_in_{db.database}"])

    # 5. 使用游标上下文（更简洁的方式）
    print("\n使用游标上下文:")
    with db.cursor_context() as cursor:
        cursor.execute("SELECT VERSION() AS version")
        version_info = cursor.fetchone()
        print(f"MySQL版本: {version_info['version']}")

    # 6. 断开连接
    if db.is_connected():
        print("\n断开数据库连接...")
        db.disconnect()

    print("\n连接状态:", "已连接" if db.is_connected() else "已断开")